import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sfepy.discrete.fem import Mesh, FEDomain, Field

def get_files_by_extension(folder_path, extension):
    """获取文件夹中指定扩展名的所有文件"""
    files = []
    for file in os.listdir(folder_path):
        if file.endswith(extension):
            files.append(file)
    return files

def shape_plot(domain: FEDomain, vecs, mode_num, freqs):
    for n_mode in range(mode_num):
        # 获取网格节点坐标
        coors = domain.mesh.coors

        # 创建图形
        plt.figure(figsize=(18, 12))

        # 绘制第n_mode个振型
        bbox = domain.get_mesh_bounding_box()
        min_coor, max_coor = bbox[:, 1]
        amp = np.abs(max_coor - min_coor) * 0.33
        # 当前振型的位移向量 (每个节点有3个分量)
        disp = vecs[:, n_mode].reshape(-1, 3)
        disp = disp / np.max(np.abs(disp)) * amp  # 将位移放大到易于观察的范围
        disp_mag = np.linalg.norm(disp, axis=1)  # 位移大小

        # 变形后的坐标 (放大位移量便于观察)
        deformed = coors + disp

        # 创建子图
        ax = plt.subplot(1, 1, 1, projection='3d')

        # 绘制原始网格点 (半透明灰色)
        ax.scatter(coors[:, 0], coors[:, 1], coors[:, 2], 
                    c='gray', alpha=0.3, s=5, label='Original')

        # 绘制变形后的网格点 (彩色显示位移大小)
        sc = ax.scatter(deformed[:, 0], deformed[:, 1], deformed[:, 2], 
                        c=disp_mag, cmap='jet', s=15, label='Deformed')

        # 添加从原始位置到变形位置的连线
        for j in range(0, len(coors), 10):  # 每隔10个点画一条线，避免太密集
            ax.plot([coors[j, 0], deformed[j, 0]],
                    [coors[j, 1], deformed[j, 1]],
                    [coors[j, 2], deformed[j, 2]],
                    'k-', linewidth=0.5, alpha=0.3)

        ax.set_title(f'Mode {n_mode+1}\nFreq: {freqs[n_mode]:.2f} Hz')
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.axis('equal')
        plt.colorbar(sc, ax=ax, label='Displacement Magnitude (mm)')
        plt.tight_layout()
        plt.suptitle('Blade Vibration Modes', y=1.02, fontsize=16)
        plt.show()

def saveas_excel(freqs, shape_vec):
    vecs_pd = pd.DataFrame(shape_vec)
    vecs_pd.columns = freqs
    vecs_pd.to_excel('modal_result.xlsx', index=False)